import math
import os
import sys
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast
from typing import Optional, Tuple, List, Union

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from plugins.models.base_model import BaseModelPlugin


class PoXiaoConfig(PretrainedConfig):
    model_type = "poxiao"

    def __init__(
            self,
            dropout: float = 0.0,
            bos_token_id: int = 1,
            eos_token_id: int = 2,
            hidden_act: str = 'silu',
            hidden_size: int = 512,
            intermediate_size: int = None,
            max_position_embeddings: int = 32768,
            num_attention_heads: int = 8,
            num_hidden_layers: int = 8,
            num_key_value_heads: int = 2,
            vocab_size: int = 6400,
            rms_norm_eps: float = 1e-05,
            rope_theta: int = 1000000.0,
            flash_attn: bool = True,

            ####################################################
            #  MOE config
            ####################################################
            use_moe: bool = False,
            num_experts_per_tok: int = 2,
            n_routed_experts: int = 4,
            n_shared_experts: int = 1,
            scoring_func: str = 'softmax',
            aux_loss_alpha: float = 0.1,
            seq_aux: bool = True,
            norm_topk_prob: bool = True,
            top_k: int = 2, 
            **kwargs
    ):
        super().__init__(**kwargs)
        self.dropout = dropout
        self.bos_token_id = bos_token_id
        self.eos_token_id = eos_token_id
        self.hidden_act = hidden_act
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.max_position_embeddings = max_position_embeddings
        self.num_attention_heads = num_attention_heads
        self.num_hidden_layers = num_hidden_layers
        self.num_key_value_heads = num_key_value_heads
        self.vocab_size = vocab_size
        self.rms_norm_eps = rms_norm_eps
        self.rope_theta = rope_theta
        self.flash_attn = flash_attn
        ####################################################
        #  MOE config
        ####################################################
        self.use_moe = use_moe
        self.num_experts_per_tok = num_experts_per_tok  # 每个token选择的专家数量
        self.n_routed_experts = n_routed_experts  # 总的专家数量
        self.n_shared_experts = n_shared_experts  # 共享专家
        self.scoring_func = scoring_func  # 评分函数，默认为'softmax'
        self.aux_loss_alpha = aux_loss_alpha  # 辅助损失的alpha参数
        self.seq_aux = seq_aux  # 是否在序列级别上计算辅助损失
        self.norm_topk_prob = norm_topk_prob  # 是否标准化top-k概率
        self.top_k = top_k


class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        """
        初始化RMSNorm层
        
        Args:
            dim (int): 输入特征的维度大小
            eps (float, optional): 用于数值稳定性的小常数. Defaults to 1e-5.
        """
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def _norm(self, x):
        """
        对输入张量进行RMS归一化操作
        
        Args:
            x (torch.Tensor): 输入张量
            
        Returns:
            torch.Tensor: 归一化后的张量
        """
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        """
        前向传播函数，对输入张量应用RMS归一化和可学习权重
        
        Args:
            x (torch.Tensor): 输入张量
            
        Returns:
            torch.Tensor: 应用RMS归一化和权重缩放后的张量
        """
        return self.weight * self._norm(x.float()).type_as(x)


def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
    """
    预计算用于旋转位置编码的频率和复数表示
    
    Args:
        dim (int): 嵌入维度
        end (int, optional): 序列最大长度. Defaults to 32768
        theta (float, optional): 频率计算参数. Defaults to 1e6
    
    Returns:
        tuple: 包含频率余弦值和正弦值的元组
    """
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device)
    freqs = torch.outer(t, freqs).float()
    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1)
    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1)
    return freqs_cos, freqs_sin


def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    """
    应用旋转位置编码到查询和键向量
    
    Args:
        q (torch.Tensor): 查询向量
        k (torch.Tensor): 键向量
        cos (torch.Tensor): 预计算的余弦值
        sin (torch.Tensor): 预计算的正弦值
        position_ids (torch.Tensor, optional): 位置索引. Defaults to None
        unsqueeze_dim (int, optional): 扩展维度. Defaults to 1
    
    Returns:
        tuple: 应用旋转位置编码后的查询和键向量
    """
    def rotate_half(x):
        """
        将张量的后半部分移到前半部分之前
        
        Args:
            x (torch.Tensor): 输入张量
            
        Returns:
            torch.Tensor: 旋转后的张量
        """
        return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)

    q_embed = (q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))
    k_embed = (k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))
    return q_embed, k_embed


def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    重复键值张量以匹配查询头的数量
    
    Args:
        x (torch.Tensor): 输入的键值张量，形状为(batch_size, sequence_length, num_key_value_heads, head_dim)
        n_rep (int): 重复次数
        
    Returns:
        torch.Tensor: 重复后的张量，形状为(batch_size, sequence_length, num_key_value_heads * n_rep, head_dim)
    """
    bs, slen, num_key_value_heads, head_dim = x.shape
    if n_rep == 1:
        return x
    return (
        x[:, :, :, None, :]
        .expand(bs, slen, num_key_value_heads, n_rep, head_dim)
        .reshape(bs, slen, num_key_value_heads * n_rep, head_dim)
    )


class MoEGate(nn.Module):
    def __init__(self, config: PoXiaoConfig):
        super().__init__()
        self.config = config
        self.top_k = config.num_experts_per_tok
        self.n_routed_experts = config.n_routed_experts

        self.scoring_func = config.scoring_func
        self.alpha = config.aux_loss_alpha
        self.seq_aux = config.seq_aux

        self.norm_topk_prob = config.norm_topk_prob
        self.gating_dim = config.hidden_size
        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
        self.reset_parameters()

    def reset_parameters(self) -> None:
        import torch.nn.init as init
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))

    def forward(self, hidden_states):
        bsz, seq_len, h = hidden_states.shape
        hidden_states = hidden_states.view(-1, h)
        logits = F.linear(hidden_states, self.weight, None)
        if self.scoring_func == 'softmax':
            scores = logits.softmax(dim=-1)
        else:
            raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')

        topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)

        if self.top_k > 1 and self.norm_topk_prob:
            denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
            topk_weight = topk_weight / denominator

        if self.training and self.alpha > 0.0:
            scores_for_aux = scores
            aux_topk = self.top_k
            topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
            if self.seq_aux:
                scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
                ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
                ce.scatter_add_(1, topk_idx_for_aux_loss,
                                torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(
                    seq_len * aux_topk / self.n_routed_experts)
                aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
            else:
                mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
                ce = mask_ce.float().mean(0)
                Pi = scores_for_aux.mean(0)
                fi = ce * self.n_routed_experts
                aux_loss = (Pi * fi).sum() * self.alpha
        else:
            aux_loss = 0
        return topk_idx, topk_weight, aux_loss


class MOEFeedForward(nn.Module):
    def __init__(self, config: PoXiaoConfig):
        super().__init__()
        self.config = config
        self.experts = nn.ModuleList([
            FeedForward(config)
            for _ in range(config.n_routed_experts)
        ])
        self.gate = MoEGate(config)
        if config.n_shared_experts > 0:
            self.shared_experts = nn.ModuleList([
                FeedForward(config)
                for _ in range(config.n_shared_experts)
            ])

    def forward(self, x):
        """
        MoE前馈网络的前向传播函数
        
        Args:
            x (torch.Tensor): 输入张量，形状为(batch_size, sequence_length, hidden_size)
            
        Returns:
            torch.Tensor: 经过MoE处理后的输出张量，形状与输入相同
        """
        identity = x
        orig_shape = x.shape
        bsz, seq_len, _ = x.shape
        # 使用门控机制选择专家
        topk_idx, topk_weight, aux_loss = self.gate(x)
        x = x.view(-1, x.shape[-1])
        flat_topk_idx = topk_idx.view(-1)
        if self.training:
            # 训练模式下，对输入进行重复并分配给不同专家处理
            x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
            y = torch.empty_like(x, dtype=torch.float16)
            for i, expert in enumerate(self.experts):
                y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype)  # 确保类型一致
            y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
            y = y.view(*orig_shape)
        else:
            # 推理模式下，使用专用的推理函数处理
            y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
        # 如果配置了共享专家，则将共享专家的输出加到结果中
        if self.config.n_shared_experts > 0:
            for expert in self.shared_experts:
                y = y + expert(identity)
        self.aux_loss = aux_loss
        return y

    @torch.no_grad()
    def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
        expert_cache = torch.zeros_like(x)
        idxs = flat_expert_indices.argsort()
        tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
        token_idxs = idxs // self.config.num_experts_per_tok
        # 当tokens_per_expert = [6, 15, 20, 26]，tokens_per_expert.shape[0]即为专家数量（此时为4）
        # 且token_idxs = [3, 7, 19, 21, 24, 25,  4,  5,  6, 10, 11, 12...] 时
        # 意味token_idxs[:6] -> [3, 7, 19, 21, 24, 25]这6个位置属于专家0处理的token（每个token有可能被多个专家处理，这取决于num_experts_per_tok）
        # 接下来9个位置token_idxs[6:15] -> [4,  5,  6, 10, 11, 12...]属于专家1处理的token...依此类推
        for i, end_idx in enumerate(tokens_per_expert):
            start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
            if start_idx == end_idx:
                continue
            expert = self.experts[i]
            exp_token_idx = token_idxs[start_idx:end_idx]
            expert_tokens = x[exp_token_idx]
            expert_out = expert(expert_tokens).to(expert_cache.dtype)
            expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
            expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out)

        return expert_cache


class FeedForward(nn.Module):
    def __init__(self, config: PoXiaoConfig):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        # 与MiniMind保持一致的中间层大小计算方式
        if config.intermediate_size is None:
            intermediate_size = int(config.hidden_size * 8 / 3)
            config.intermediate_size = 64 * ((intermediate_size + 64 - 1) // 64)
        self.ffn_dim = config.intermediate_size
        
        self.gate_proj = nn.Linear(self.hidden_size, self.ffn_dim, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.ffn_dim, bias=False)
        self.down_proj = nn.Linear(self.ffn_dim, self.hidden_size, bias=False)
        self.dropout = nn.Dropout(config.dropout)
        # 使用与MiniMind相同的激活函数
        self.act_fn = nn.SiLU() if config.hidden_act == 'silu' else nn.GELU()

    def forward(self, hidden_states):
        # 与MiniMind保持一致的前向传播计算方式
        return self.dropout(self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)))


class Attention(nn.Module):
    def __init__(self, config: PoXiaoConfig):
        super().__init__()
        self.config = config
        self.num_key_value_heads = config.num_key_value_heads if config.num_key_value_heads is not None else config.num_attention_heads
        assert config.num_attention_heads % self.num_key_value_heads == 0
        self.n_local_heads = config.num_attention_heads
        self.n_local_kv_heads = self.num_key_value_heads
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        self.head_dim = config.hidden_size // config.num_attention_heads
        
        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
        
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.dropout = config.dropout
        
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and config.flash_attn

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        use_cache: bool = False,
    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
        bsz, q_len, _ = hidden_states.size()

        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)

        query_states = query_states.view(bsz, q_len, self.n_local_heads, self.head_dim)
        key_states = key_states.view(bsz, q_len, self.n_local_kv_heads, self.head_dim)
        value_states = value_states.view(bsz, q_len, self.n_local_kv_heads, self.head_dim)

        # 应用旋转位置编码
        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        if past_key_value is not None:
            # 重用缓存的key和value
            key_states = torch.cat([past_key_value[0], key_states], dim=1)
            value_states = torch.cat([past_key_value[1], value_states], dim=1)

        past_key_value = (key_states, value_states) if use_cache else None

        # 重复KV以匹配查询头数
        key_states = repeat_kv(key_states, self.n_rep)
        value_states = repeat_kv(value_states, self.n_rep)

        # 转换形状以适应注意力计算
        query_states = query_states.transpose(1, 2)  # (bsz, n_local_heads, q_len, head_dim)
        key_states = key_states.transpose(1, 2)      # (bsz, n_local_heads, kv_len, head_dim)
        value_states = value_states.transpose(1, 2)  # (bsz, n_local_heads, kv_len, head_dim)

        if self.flash and q_len != 1:
            # 使用Flash Attention
            dropout_p = self.dropout if self.training else 0.0
            attn_mask = None
            if attention_mask is not None:
                attn_mask = attention_mask.view(bsz, 1, 1, -1).expand(bsz, self.n_local_heads, q_len, -1)
                attn_mask = attn_mask.to(dtype=query_states.dtype)  # 确保数据类型一致

            attn_output = torch.nn.functional.scaled_dot_product_attention(
                query_states, key_states, value_states, 
                attn_mask=attn_mask, 
                dropout_p=dropout_p,
                is_causal=True  # 添加因果掩码，与MiniMind保持一致
            )
        else:
            # 标准注意力计算，添加因果掩码
            scores = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
            
            # 添加因果掩码，与MiniMind保持一致
            scores = scores + torch.triu(
                torch.full((q_len, q_len), float("-inf"), device=scores.device),
                diagonal=1
            ).unsqueeze(0).unsqueeze(0)
            
            if attention_mask is not None:
                # 处理外部传入的attention_mask
                # 修复bool tensor减法错误，与MiniMind保持一致
                extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
                extended_attention_mask = extended_attention_mask.to(dtype=scores.dtype)  # fp16 compatibility
                extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(scores.dtype).min
                scores = scores + extended_attention_mask

            scores = F.softmax(scores.float(), dim=-1).type_as(query_states)
            scores = self.attn_dropout(scores)
            attn_output = torch.matmul(scores, value_states)

        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.reshape(bsz, q_len, self.n_local_heads * self.head_dim)
        attn_output = self.o_proj(attn_output)
        attn_output = self.resid_dropout(attn_output)

        return attn_output, past_key_value


class DecoderLayer(nn.Module):
    def __init__(self, config: PoXiaoConfig):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        
        self.self_attn = Attention(config)
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        
        if config.use_moe:
            self.mlp = MOEFeedForward(config)
        else:
            self.mlp = FeedForward(config)
            
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        use_cache: bool = False,
    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
        residual = hidden_states
        
        # 自注意力
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            position_embeddings=position_embeddings,
            attention_mask=attention_mask,
            past_key_value=past_key_value,
            use_cache=use_cache,
        )
        hidden_states = residual + hidden_states

        # 前馈网络
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        if self.config.use_moe:
            hidden_states = self.mlp(hidden_states)
            aux_loss = self.mlp.aux_loss
        else:
            hidden_states = self.mlp(hidden_states)
            aux_loss = 0
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)
        if use_cache:
            outputs += (present_key_value,)
            
        outputs += (aux_loss,)
        return outputs


class PoXiaoModel(nn.Module):
    def __init__(self, config: PoXiaoConfig):
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_hidden_layers)])
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.dropout = nn.Dropout(config.dropout)
        
        # 预计算位置编码
        freqs_cos, freqs_sin = precompute_freqs_cis(
            config.hidden_size // config.num_attention_heads, 
            config.max_position_embeddings,
            config.rope_theta
        )
        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
        self.register_buffer("freqs_sin", freqs_sin, persistent=False)
        
    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
        batch_size, seq_length = input_ids.shape
        past_key_values = past_key_values or [None] * len(self.layers)
        start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0

        hidden_states = self.dropout(self.embed_tokens(input_ids))

        position_embeddings = (
            self.freqs_cos[start_pos:start_pos + seq_length],
            self.freqs_sin[start_pos:start_pos + seq_length]
        )

        # 逐层处理输入
        presents = []
        aux_loss = 0
        for layer_idx, (layer, past_key_value) in enumerate(zip(self.layers, past_key_values)):
            layer_outputs = layer(
                hidden_states,
                position_embeddings,
                attention_mask=attention_mask,
                past_key_value=past_key_value,
                use_cache=use_cache,
            )
            hidden_states = layer_outputs[0]
            
            if use_cache:
                presents.append(layer_outputs[1])
                
            if self.config.use_moe:
                aux_loss += layer_outputs[2]

        hidden_states = self.norm(hidden_states)
        return hidden_states, presents, aux_loss


class PoXiaoForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = PoXiaoConfig

    def __init__(self, config: PoXiaoConfig = None):
        """
        初始化因果语言模型
        
        Args:
            config (PoXiaoConfig, optional): 模型配置对象. 如果未提供则使用默认配置.
        """
        super().__init__(config)
        self.model = PoXiaoModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.model.embed_tokens.weight = self.lm_head.weight

    def forward(self,
                input_ids: Optional[torch.Tensor] = None,
                attention_mask: Optional[torch.Tensor] = None,
                past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
                use_cache: bool = False,
                labels: Optional[torch.Tensor] = None,
                **kwargs):
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            **kwargs
        )
        
        hidden_states = outputs[0]
        logits = self.lm_head(hidden_states)
        
        loss = None
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = nn.CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # Enable model parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)
            
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs[1],
        )

    def generate(self, input_ids, max_length, **kwargs):
        """生成文本"""
        # 使用transformers库的GenerationMixin.generate方法
        return super().generate(input_ids, max_length=max_length, **kwargs)


class PoXiaoModelPlugin(BaseModelPlugin):
    """PoXiao模型插件"""
    
    def create_model(self, config):
        """创建PoXiao模型实例"""
        model_config = PoXiaoConfig(
            hidden_size=config.get('hidden_size', 512),
            num_hidden_layers=config.get('num_hidden_layers', 8),
            use_moe=config.get('use_moe', False),
            # 与MiniMind保持一致的配置
            intermediate_size=None,  # 使用与MiniMind相同的计算方式
            hidden_act='silu'  # 使用SiLU激活函数
        )
        model = PoXiaoForCausalLM(model_config)
        
        # 如果配置中启用了Kaiming初始化，则对模型参数进行初始化
        if config.get('use_kaiming_init', False):
            self._init_weights_kaiming(model)
            
        return model
        
    def get_model_params_count(self, model):
        """获取模型参数量"""
        return sum(p.numel() for p in model.parameters() if p.requires_grad)
        
    def _init_weights_kaiming(self, model):
        """
        使用与MiniMind一致的初始化方法初始化模型权重
        参考litgpt/extensions/thunder/pretrain.py中的GPT-NeoX权重初始化
        
        Args:
            model (nn.Module): 需要初始化的模型
        """
        def init_weights(module, std):
            if hasattr(module, 'weight') and module.weight is not None:
                torch.nn.init.normal_(module.weight, mean=0.0, std=std)
            if hasattr(module, 'bias') and module.bias is not None:
                torch.nn.init.zeros_(module.bias)
                
        n_embd = self.config.hidden_size if hasattr(self, 'config') else 512
        n_layer = self.config.num_hidden_layers if hasattr(self, 'config') else 8
        
        std = math.sqrt(2.0 / 5 / n_embd)
        
        for mod in model.modules():
            if isinstance(mod, (nn.Embedding, nn.Linear)):
                init_weights(mod, std=std)
            elif isinstance(mod, (nn.LayerNorm, RMSNorm)):
                # 对归一化层的权重初始化为1，偏置初始化为0
                if hasattr(mod, 'weight') and mod.weight is not None:
                    torch.nn.init.ones_(mod.weight)
                if hasattr(mod, 'bias') and mod.bias is not None:
                    torch.nn.init.zeros_(mod.bias)